import asyncio
import logging
from abc import ABCMeta, abstractmethod
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Dict, Optional

from hamilton.async_driver import AsyncDriver
from hamilton.driver import Driver
from haystack import Pipeline

from src.core.engine import Engine
from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider

logger = logging.getLogger("wren-ai-service")

class BasicPipeline(metaclass=ABCMeta):
    def __init__(self, pipe: Pipeline | AsyncDriver | Driver):
        self._pipe = pipe

    @abstractmethod
    def run(self, *args, **kwargs) -> Dict[str, Any]:
        ...


def async_validate(task: callable):
    result = asyncio.run(task())
    print(result)
    return result


@dataclass
class PipelineComponent(Mapping):
    llm_provider: LLMProvider = None
    embedder_provider: EmbedderProvider = None
    document_store_provider: DocumentStoreProvider = None
    engine: Engine = None

    def __getitem__(self, key):
        return getattr(self, key)

    def __iter__(self):
        return iter(self.__dict__)

    def __len__(self):
        return len(self.__dict__)


class PipelineManager:
    def __init__(self, pipelines: Dict[str, BasicPipeline]):
        self._pipelines = pipelines
        self._results = {}
        self._background_tasks = None
        self.id = None

    def set_background_tasks(self, background_tasks):
        self._background_tasks = background_tasks

    def set_id(self, id):
        self.id = id

    async def run_async(self, pipeline_name: str, **kwargs):
        """Run a pipeline asynchronously."""
        if pipeline_name not in self._pipelines:
            logger.error(f"Pipeline {pipeline_name} not found")
            raise ValueError(f"Pipeline {pipeline_name} not found")

        logger.info(f"Running pipeline {pipeline_name} asynchronously with ID {self.id}")
        self._results[self.id] = None
        self._background_tasks.add_task(self._run_pipeline, pipeline_name, **kwargs)

    async def _run_pipeline(self, pipeline_name: str, **kwargs):
        """Run a pipeline and store the result."""
        try:
            logger.info(f"Starting pipeline {pipeline_name} with ID {self.id}")
            result = await self._pipelines[pipeline_name].run(**kwargs)
            
            # Log the result for debugging
            logger.info(f"Pipeline {pipeline_name} result for ID {self.id}: {result}")
            
            # Process the result based on pipeline type
            if pipeline_name == "question_recommendation":
                # For question_recommendation, ensure we have a list of questions
                processed_result = self._process_question_recommendation_result(result)
                self._results[self.id] = processed_result
                logger.info(f"Processed question_recommendation result for ID {self.id}: {processed_result}")
            else:
                self._results[self.id] = result
                
            logger.info(f"Pipeline {pipeline_name} completed for ID {self.id}")
        except Exception as e:
            logger.error(f"An error occurred during {pipeline_name} generation for ID {self.id}: {e}")
            # Store an empty result instead of None to indicate completion with error
            if pipeline_name == "question_recommendation":
                self._results[self.id] = []  # Empty list for question_recommendation
            else:
                self._results[self.id] = {}  # Empty dict for other pipelines

    def _process_question_recommendation_result(self, result):
        """Process the result from question_recommendation pipeline."""
        try:
            logger.info(f"Processing question_recommendation result: {result}")
            
            # If result is None, return an empty list
            if result is None:
                logger.warning("Result is None, returning empty list")
                return []
            
            # Handle different result formats
            if isinstance(result, list):
                # Result is already a list of questions
                logger.info("Result is a list")
                # Ensure each item has question and category
                processed_questions = []
                for item in result:
                    if isinstance(item, dict) and "question" in item and "category" in item:
                        processed_questions.append(item)
                    else:
                        logger.warning(f"Invalid question format: {item}")
                return processed_questions
            elif isinstance(result, dict):
                # If result is a dict, try to get normalized or questions
                logger.info("Result is a dict")
                if "normalized" in result:
                    normalized = result["normalized"]
                    logger.info(f"Found normalized in result: {normalized}")
                    if isinstance(normalized, list):
                        # Ensure each item has question and category
                        processed_questions = []
                        for item in normalized:
                            if isinstance(item, dict) and "question" in item and "category" in item:
                                processed_questions.append(item)
                            else:
                                logger.warning(f"Invalid question format in normalized: {item}")
                        return processed_questions
                    elif isinstance(normalized, dict) and "questions" in normalized:
                        questions = normalized["questions"]
                        if isinstance(questions, list):
                            # Ensure each item has question and category
                            processed_questions = []
                            for item in questions:
                                if isinstance(item, dict) and "question" in item and "category" in item:
                                    processed_questions.append(item)
                                else:
                                    logger.warning(f"Invalid question format in normalized.questions: {item}")
                            return processed_questions
                elif "questions" in result:
                    questions = result["questions"]
                    logger.info(f"Found questions in result: {questions}")
                    # If questions is a dict (category -> list of questions), flatten it
                    if isinstance(questions, dict):
                        logger.info("Questions is a dict, flattening")
                        flattened = []
                        for category, category_questions in questions.items():
                            if isinstance(category_questions, list):
                                for q in category_questions:
                                    if isinstance(q, dict) and "question" in q:
                                        flattened.append({"question": q["question"], "category": category})
                                    else:
                                        logger.warning(f"Invalid question format in questions dict: {q}")
                        return flattened
                    elif isinstance(questions, list):
                        # Ensure each item has question and category
                        processed_questions = []
                        for item in questions:
                            if isinstance(item, dict) and "question" in item and "category" in item:
                                processed_questions.append(item)
                            else:
                                logger.warning(f"Invalid question format in questions list: {item}")
                        return processed_questions
            # Default to empty list if we can't determine the format
            logger.error(f"Unexpected result format: {result}")
            return []
        except Exception as e:
            logger.error(f"Error processing question recommendation result: {e}")
            return []

    async def get_result(self, id: str):
        """Get the result of a pipeline run."""
        result = self._results.get(id)
        logger.info(f"Getting result for {id}: {result}")
        return result

    async def run(self, pipeline_name: str, **kwargs):
        """Run a pipeline synchronously."""
        if pipeline_name not in self._pipelines:
            raise ValueError(f"Pipeline {pipeline_name} not found")

        return await self._pipelines[pipeline_name].run(**kwargs)
